#!/usr/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import os
import sys

# 添加项目根目录到Python搜索路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from vector_store.factory import VectorStoreFactory
from vector_store.node_vector_processor import NodeVectorProcessor
from config import config

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

def main():
    """主入口函数"""
    parser = argparse.ArgumentParser(description='向量存储查询工具')
    parser.add_argument('--query', type=str, help='搜索查询文本')
    parser.add_argument('--store-type', type=str, help='向量存储类型 (sqlite 或 milvus)')
    parser.add_argument('--sqlite-db-path', type=str, help='SQLite数据库路径')
    parser.add_argument('--milvus-db-path', type=str, help='Milvus本地数据库路径')
    parser.add_argument('--milvus-uri', type=str, help='Milvus服务器URI地址')
    parser.add_argument('--milvus-token', type=str, help='Milvus服务器访问令牌')
    parser.add_argument('--top-k', type=int, default=5, help='返回结果数量')
    parser.add_argument('--node-type', type=str, choices=['class', 'function', 'annotation'], help='节点类型过滤')
    parser.add_argument('--api-key', type=str, help='嵌入API密钥')
    parser.add_argument('--log-level', type=str, default='INFO', 
                       choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'], help='日志级别')
    
    args = parser.parse_args()
    
    # 设置日志级别
    logging.getLogger().setLevel(getattr(logging, args.log_level))
    
    if not args.query:
        parser.print_help()
        return
    
    try:
        # 更新环境变量配置
        if args.store_type:
            os.environ["VECTOR_STORE_TYPE"] = args.store_type
        if args.sqlite_db_path:
            os.environ["SQLITE_DB_PATH"] = args.sqlite_db_path
        if args.milvus_db_path:
            os.environ["MILVUS_DB_PATH"] = args.milvus_db_path
        if args.milvus_uri:
            os.environ["MILVUS_URI"] = args.milvus_uri
        if args.milvus_token:
            os.environ["MILVUS_TOKEN"] = args.milvus_token
        if args.api_key:
            os.environ["EMBEDDING_API_KEY"] = args.api_key
            
        # 重新加载配置
        from config.config import Config
        global config
        config = Config()
        
        # 使用配置创建向量存储
        vector_store = VectorStoreFactory.create_vector_store()
        
        # 创建节点向量处理器
        processor = NodeVectorProcessor(vector_store=vector_store)
        
        # 搜索相似节点
        results = processor.search_similar_nodes(
            query_text=args.query,
            node_type=args.node_type,
            top_k=args.top_k
        )
        
        # 显示结果
        if not results:
            print(f"未找到与查询'{args.query}'相关的结果")
            return
        
        print(f"\n找到 {len(results)} 个与查询'{args.query}'相关的结果:\n")
        
        for i, result in enumerate(results, 1):
            print(f"结果 {i}:")
            print(f"  节点类型: {result['node_type']}")
            print(f"  原始ID: {result['original_id']}")
            print(f"  相似度: {result['similarity']:.4f}")
            
            metadata = result.get('metadata', {})
            if metadata:
                if 'name' in metadata:
                    print(f"  名称: {metadata['name']}")
                if 'full_name' in metadata:
                    print(f"  完整名称: {metadata['full_name']}")
                if 'line_start' in metadata and 'line_end' in metadata:
                    print(f"  行范围: {metadata['line_start']}-{metadata['line_end']}")
            
            print()
        
    except Exception as e:
        logger.error(f"发生错误: {str(e)}")
        return 1
    finally:
        if 'processor' in locals():
            processor.close()
    
    return 0

if __name__ == "__main__":
    sys.exit(main()) 